#! -*- coding:utf-8 -*-
from asyncio import FastChildWatcher
from tensorflow.python import keras
#from tensorflow.python.keras.optimizers import adam_v2
from keras.optimizers import Adam
from data_loader import data_generator, load_data
from model import E2EModel, Eval
from utils import extract_items, get_tokenizer, metric
import os, argparse
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
from keras import backend as K
import tensorflow as tf

parser = argparse.ArgumentParser(description='Model Controller')
parser.add_argument('--train', default=True, type=bool, help='to train the HBT model, python run.py --train=True')
parser.add_argument('--dataset', default='WebNLG', type=str, help='specify the dataset from ["NYT","WebNLG"')
args = parser.parse_args()


if __name__ == '__main__':
    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
    bert_model = 'cased_L-12_H-768_A-12'
    bert_config_path = './pretrained_bert_models/' + bert_model + '/bert_config.json'
    bert_vocab_path = './pretrained_bert_models/' + bert_model + '/vocab.txt'
    bert_checkpoint_path = './pretrained_bert_models/' + bert_model + '/bert_model.ckpt'
    dataset = args.dataset
    train_path = './data/' + dataset + '/new/train_triples.json'
    dev_path = './data/' + dataset + '/dev_triples.json'
    #test_path = './data/' + dataset + '/test_split_by_num/test_triples_5.json' # ['1','2','3','4','5']
    #test_path = './data/' + dataset + '/test_split_by_type/test_triples_normal.json' # ['normal', 'seo', 'epo']
    test_path = './data/' + dataset + '/test_triples.json' # overall test
    rel_dict_path = './data/' + dataset + '/rel2id.json'
    save_weights_path = './saved_weights/' + dataset + '/best_model.weights'
    dir_dict_path =  './data/' + dataset + '/dir2id.json'
    load_path = './saved_weights/' + dataset + '/backhold/best_model.weights'
    LR = 1e-5
    tokenizer = get_tokenizer(bert_vocab_path)#
    train_data, dev_data, test_data, id2rel, rel2id, num_rels, id2dir, dir2id, num_dirs = load_data(train_path, dev_path, test_path, rel_dict_path, dir_dict_path)
   
    with strategy.scope():

        entity_model, relation_model, direction_model, joint_model = E2EModel(bert_config_path, bert_checkpoint_path, LR, num_rels, num_dirs)
        joint_model.compile(optimizer=Adam(LR))
    
    
    
    if args.train:
        print('>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>train begin!>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>')
        BATCH_SIZE = 4
        EPOCH = 9999
        MAX_LEN = 100
        STEPS = len(train_data) // BATCH_SIZE 
        data_manager = data_generator(train_data, tokenizer, rel2id, id2rel, num_rels, dir2id, id2dir, num_dirs, MAX_LEN, BATCH_SIZE)
        evaluator = Eval(entity_model, relation_model,direction_model, tokenizer, id2rel, dev_data, save_weights_path, id2dir)
        #joint_model.load_weights(save_weights_path)
        joint_model.fit_generator(data_manager.__iter__(),
                              steps_per_epoch= STEPS,
                              epochs= EPOCH,
                              callbacks= [evaluator],
                              )
    else:
        join_model.load_weights(save_weights_path)
        test_result_path = './results/' + dataset + '/test_result_new.json'
        isExactMatch = True if dataset == 'NYT11' else False
        if isExactMatch:
            print("Exact Match")
        else:
            print("Partial Match")
        e_precision, e_recall, e_f1, r_precision, r_recall, r_f1, triples_precision, triples_recall, triples_f1,triple_correct_num, triple_predict_num, triple_gold_num,dir4_num, fake_num = metric(entity_model, relation_model, direction_model,  test_data, id2rel ,tokenizer, id2dir,isExactMatch, test_result_path)
        print(f'{e_precision}\t{e_recall}\t{e_f1}\t{r_precision}\t{r_recall}\t{r_f1}\t{triples_precision}\t{triples_recall}\t{triples_f1}')
        print(dir4_num)
        print(fake_num)
        print("Test is over!")
